from gensim import corpora
from handler.SentenceHandler import SentenceHandler
from collections import defaultdict

class DataPreHandler:

    def __init__(self, data):
        self.simpleCorpus = None
        self.dictionary = None
        self.sentences = []

        for i in range(0, len(data)):
            sentence = SentenceHandler(data[i], i)
            self.sentences.append(sentence)

    def preModel(self, min_frequency = 1):
        texts = []
        for sentence in self.sentences:
            texts.append(sentence.get_cut_sentence())

        # 删除低频词
        frequency = defaultdict(int)
        for text in texts:
            for token in text:
                frequency[token] += 1

        texts = [[token for token in text if frequency[token] > min_frequency] for text in texts]

        self.dictionary = corpora.Dictionary(texts)
        self.simpleCorpus = [self.dictionary.doc2bow(text) for text in texts]

